import argparse
import xml.etree.ElementTree as ET
from pathlib import Path

import matplotlib.pyplot as plt
import seaborn as sns
from mmcv import Config
from tqdm import tqdm


def argument_parser():
    parser = argparse.ArgumentParser(description="dataset test")
    parser.add_argument("--config-file",
                        default="",
                        metavar="FILE",
                        help="path to config file")
    args = parser.parse_args()
    return args


def get_box_size_list(xml_file_path_list):
    box_size_list = list()
    for xml_file_path in tqdm(xml_file_path_list, ascii=True):
        tree = ET.parse(xml_file_path)
        root = tree.getroot()
        for obj in root.findall("object"):
            name = obj.find("name").text
            if name not in ("person", "head"):
                continue
            bnd_box = obj.find("bndbox")
            bbox = [
                int(float(bnd_box.find("xmin").text)),
                int(float(bnd_box.find("ymin").text)),
                int(float(bnd_box.find("xmax").text)),
                int(float(bnd_box.find("ymax").text))
            ]
            size = (bbox[2] - bbox[0], bbox[3] - bbox[1])
            box_size_list.append(size)
    return box_size_list


def draw_size_distribution(box_size_list):
    x, y = list(zip(*box_size_list))
    sns.set_theme()
    sns.jointplot(x=x, y=y)
    plt.xlabel("width")
    plt.ylabel("height")
    plt.title("size distribution")
    plt.show()


def main():
    args = argument_parser()
    config_file = args.config_file
    cfg = Config.fromfile(config_file)
    data_root = cfg.data_root
    xml_file_path_list = [str(path) for path in Path(data_root).rglob("*.xml")]
    box_size_list = get_box_size_list(xml_file_path_list)
    draw_size_distribution(box_size_list)


if __name__ == "__main__":
    main()
